{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c4ed29b-f19c-4b28-8cd1-cf19d793e298",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ============================================================\n",
    "# Nearest-Neighbour Mapping\n",
    "# ============================================================\n",
    "import netCDF4 as nc\n",
    "import numpy as np\n",
    "from scipy.spatial import cKDTree\n",
    "from scipy.interpolate import griddata\n",
    "import time\n",
    "import os\n",
    "\n",
    "# ============================================================\n",
    "# 1. File\n",
    "# ============================================================\n",
    "SOURCE_GRID = ''\n",
    "SOURCE_INFO = ''\n",
    "SOURCE_MEAN = ''\n",
    "TARGET_INFO = ''\n",
    "OUTPUT_MEAN = ''\n",
    "OUTPUT_VTK = ''\n",
    "\n",
    "os.makedirs(os.path.dirname(OUTPUT_MEAN), exist_ok=True)\n",
    "\n",
    "# ============================================================\n",
    "# 2. Loading original coordinates\n",
    "# ============================================================\n",
    "source_grid = nc.Dataset(SOURCE_GRID, 'r')\n",
    "source_info = nc.Dataset(SOURCE_INFO, 'r')\n",
    "source_x = source_info.variables['coord_0'][:]\n",
    "source_y = source_info.variables['coord_1'][:]\n",
    "source_coords = np.column_stack([source_x, source_y])\n",
    "\n",
    "# ============================================================\n",
    "# 3. Loading original data\n",
    "# ============================================================\n",
    "source_mean = nc.Dataset(SOURCE_MEAN, 'r')\n",
    "\n",
    "# Reading original data\n",
    "u_source = source_mean.variables['variables0'][:]\n",
    "v_source = source_mean.variables['variables1'][:]\n",
    "rho_source = source_mean.variables['variables2'][:]\n",
    "\n",
    "cell_info = source_grid.variables['cellInfo'][:]\n",
    "valid_source_mask = (cell_info == 0)\n",
    "\n",
    "source_x_valid = source_x[valid_source_mask]\n",
    "source_y_valid = source_y[valid_source_mask]\n",
    "source_coords_valid = np.column_stack([source_x_valid, source_y_valid])\n",
    "u_source_valid = u_source[valid_source_mask]\n",
    "v_source_valid = v_source[valid_source_mask]\n",
    "rho_source_valid = rho_source[valid_source_mask]\n",
    "\n",
    "# ============================================================\n",
    "# 4. Targetting\n",
    "# ============================================================\n",
    "target_info = nc.Dataset(TARGET_INFO, 'r')\n",
    "target_x = target_info.variables['coord_0'][:]\n",
    "target_y = target_info.variables['coord_1'][:]\n",
    "target_coords = np.column_stack([target_x, target_y])\n",
    "\n",
    "# ============================================================\n",
    "# 5. KDTree\n",
    "# ============================================================\n",
    "t0 = time.time()\n",
    "tree = cKDTree(source_coords_valid)\n",
    "\n",
    "# ============================================================\n",
    "# 6. Nearest-Neighbour\n",
    "# ============================================================\n",
    "t0 = time.time()\n",
    "distances, indices = tree.query(target_coords, k=1)\n",
    "\n",
    "u_mapped = u_source_valid[indices]\n",
    "v_mapped = v_source_valid[indices]\n",
    "rho_mapped = rho_source_valid[indices]\n",
    "        \n",
    "# ============================================================\n",
    "# 8. NetCDF File\n",
    "# ============================================================\n",
    "output = nc.Dataset(OUTPUT_MEAN, 'w', format='NETCDF4')\n",
    "output.createDimension('nPoints', len(target_x))\n",
    "\n",
    "var_u = output.createVariable('variables0', 'f8', ('nPoints',))\n",
    "var_v = output.createVariable('variables1', 'f8', ('nPoints',))\n",
    "var_rho = output.createVariable('variables2', 'f8', ('nPoints',))\n",
    "\n",
    "if 'globalTimeStep' in source_mean.variables.keys():\n",
    "    global_time = output.createVariable('globalTimeStep', \n",
    "                                    source_mean.variables['globalTimeStep'].dtype,\n",
    "                                    source_mean.variables['globalTimeStep'].dimensions)\n",
    "    global_time[:] = source_mean.variables['globalTimeStep'][:]\n",
    "\n",
    "var_u[:] = u_mapped\n",
    "var_v[:] = v_mapped\n",
    "var_rho[:] = rho_mapped\n",
    "\n",
    "var_u.long_name = 'um'\n",
    "var_v.long_name = 'vm'\n",
    "var_rho.long_name = 'rhom'\n",
    "\n",
    "output.description = 'Mean velocity data'\n",
    "output.source_file = SOURCE_MEAN\n",
    "output.target_grid = TARGET_INFO\n",
    "output.mapping_method = 'nearest-neighbour'\n",
    "output.n_source_points = len(source_x)\n",
    "output.n_valid_source_points = len(source_x_valid)\n",
    "output.n_target_points = len(target_x)\n",
    "output.mean_mapping_distance = float(distances.mean())\n",
    "output.max_mapping_distance = float(distances.max())\n",
    "output.creation_date = time.strftime('%Y-%m-%d %H:%M:%S')\n",
    "\n",
    "output.close()\n",
    "\n",
    "print(f\"✓ NetCDF\")\n",
    "\n",
    "# ============================================================\n",
    "# 10. ParaView VTK File\n",
    "# ============================================================\n",
    "cx = np.asarray(target_x)\n",
    "cy = np.asarray(target_y)\n",
    "\n",
    "if np.ma.is_masked(cx): cx = cx.filled()\n",
    "if np.ma.is_masked(cy): cy = cy.filled()\n",
    "\n",
    "ux = np.unique(cx)\n",
    "uy = np.unique(cy)\n",
    "\n",
    "dx = np.median(np.diff(ux))\n",
    "dy = np.median(np.diff(uy))\n",
    "\n",
    "points = []\n",
    "cells = []\n",
    "cell_u = []\n",
    "cell_v = []\n",
    "cell_rho = []\n",
    "\n",
    "point_idx = {}\n",
    "pid = 0\n",
    "\n",
    "for i in range(len(cx)):\n",
    "    x = cx[i]\n",
    "    y = cy[i]\n",
    "\n",
    "    quad = []\n",
    "    for vx, vy in [\n",
    "        (x - dx/2, y - dy/2),\n",
    "        (x + dx/2, y - dy/2),\n",
    "        (x + dx/2, y + dy/2),\n",
    "        (x - dx/2, y + dy/2),\n",
    "    ]:\n",
    "        key = (vx, vy)\n",
    "        if key not in point_idx:\n",
    "            point_idx[key] = pid\n",
    "            points.append((vx, vy, 0.0))\n",
    "            pid += 1\n",
    "        quad.append(point_idx[key])\n",
    "\n",
    "    cells.append(quad)\n",
    "\n",
    "    cell_u.append(u_mapped[i])\n",
    "    cell_v.append(v_mapped[i])\n",
    "    cell_rho.append(rho_mapped[i])\n",
    "\n",
    "# VTK\n",
    "with open(OUTPUT_VTK, 'w') as f:\n",
    "    f.write(\"# vtk DataFile Version 3.0\\n\")\n",
    "    f.write(\"Mean velocity (grid_small_coarse, cell-centered)\\n\")\n",
    "    f.write(\"ASCII\\n\")\n",
    "    f.write(\"DATASET UNSTRUCTURED_GRID\\n\")\n",
    "\n",
    "    f.write(f\"\\nPOINTS {len(points)} float\\n\")\n",
    "    for p in points:\n",
    "        f.write(f\"{p[0]:.6f} {p[1]:.6f} {p[2]:.6f}\\n\")\n",
    "\n",
    "    f.write(f\"\\nCELLS {len(cells)} {len(cells)*5}\\n\")\n",
    "    for c in cells:\n",
    "        f.write(f\"4 {c[0]} {c[1]} {c[2]} {c[3]}\\n\")\n",
    "\n",
    "    f.write(f\"\\nCELL_TYPES {len(cells)}\\n\")\n",
    "    for _ in cells:\n",
    "        f.write(\"9\\n\")\n",
    "\n",
    "    f.write(f\"\\nCELL_DATA {len(cells)}\\n\")\n",
    "    \n",
    "    f.write(\"\\nSCALARS rhom float 1\\nLOOKUP_TABLE default\\n\")\n",
    "    for v in cell_rho:\n",
    "        f.write(f\"{v:.6f}\\n\")\n",
    "\n",
    "    f.write(\"\\nSCALARS um float 1\\nLOOKUP_TABLE default\\n\")\n",
    "    for v in cell_u:\n",
    "        f.write(f\"{v:.6f}\\n\")\n",
    "\n",
    "    f.write(\"\\nSCALARS vm float 1\\nLOOKUP_TABLE default\\n\")\n",
    "    for v in cell_v:\n",
    "        f.write(f\"{v:.6f}\\n\")\n",
    "\n",
    "print(\"✓ VTK\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
